import cv2
import torch
import os
import uuid

def process_video(video_path):
    # 创建 outputs 文件夹（如果不存在）
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成唯一的输出文件名
    output_filename = f"processed_video_{uuid.uuid4().hex}.mp4"
    output_path = os.path.join(output_dir, output_filename)

    # 加载模型
    model = torch.hub.load('./yolov5', 'custom', path='./models/hats_best.pt', source='local')

    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame)
        rendered_frame = results.render()[0]

        cv2.imshow('YOLOv5 Detection', rendered_frame)
        out.write(rendered_frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    out.release()
    cv2.destroyAllWindows()